import torch

# 清除 CUDA 缓存并复位最大内存计数器
torch.cuda.reset_max_memory_allocated()
torch.cuda.empty_cache()

# 初始化一些变量
total_storage_nbytes = 0
tensors = []

# 创建多个张量
for _ in range(900):
    tensor = torch.randn(1024, 1024).cuda()
    tensors.append(tensor)
    total_storage_nbytes += tensor.untyped_storage().nbytes()

# 打印所有张量实际占用的内存总量
print(f"Total storage nbytes: {total_storage_nbytes / (1024 ** 2)} MB")

# 打印最大显存使用量
print(f"Max memory allocated: {torch.cuda.max_memory_allocated() / (1024 ** 2)} MB")

# 对比两者的差异
difference = torch.cuda.max_memory_allocated() - total_storage_nbytes
print(f"Difference: {difference / (1024 ** 2)} MB")


